load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_code.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_csm.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/contributed.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/shea_util.ncl"
;**************************************
begin

modelname = getenv("modelname")
varname = getenv("varname")
season = getenv("season")

;modelname = "ACCESS-CM2"
;varname = "pr5d"
;season = "JJA"


Nlevel = 6                 ;6 warming levels

nlat = 180
nlon = 360
threshold = 0.1            ;wet/dry events threshold: 0.1mm/day


;----------------warming periods------------------
ft := addfile("/WORK/zhangwx/EConstraint/data/CMIP6/tas/GMST/GMST_"+modelname+"_warming_levels_year.nc","r")
year_center := ft->year_center             ;(6 level) for 0,0.5,1,2,3,4C wrt baseline


;----------------pr data------------------
diri := "/WORK/zhangwx/EConstraint/data/raw_prNd/CMIP6/"+varname+"/"+season+"/"
fi := addfile(diri+varname+"_"+modelname+"_historical_ssp585_1979-2100_"+season+".nc","r")
var := fi->$varname$              ;(day lat lon)   mm/day

time := var&time
date := cd_calendar(time,0)
year := date(:,0)

;------------------------------------------
;***************************************
k_shape := new((/Nlevel,nlat,nlon/),float)
k_shape = k_shape@_FillValue
theta_scale := k_shape

mean_direct := k_shape
Var_direct := k_shape             ;direct estimates from model data (for wet days)
frq_dry := k_shape        ;frequency of dry events


do ilevel = 0,Nlevel-1

  if (ismissing(year_center(ilevel))) then
    continue
  end if

  yrstrt := year_center(ilevel)-9
  yrlast := year_center(ilevel)+10
  iselected := ind( (year .ge. yrstrt).and.(year .le. yrlast) )
  var_ilevel := var(iselected,:,:)               ;(~1800day lat lon)  all-day

  do ilat = 0,nlat-1
    do ilon = 0,nlon-1

    iwet := ind( var_ilevel(:,ilat,ilon) .ge. threshold)     ;index for wet events
    if (all(ismissing(iwet))) then
      frq_dry(ilevel,ilat,ilon) = 1.0
      continue
    end if
    var_ilevel_wet := var_ilevel(iwet,ilat,ilon)        ;(wetdays) 1dim

    frq_dry(ilevel,ilat,ilon) = 1.0-tofloat(dimsizes(iwet))/dimsizes(iselected)         ;dry-event frequency

    gamma_par := dim_gamfit_n(var_ilevel_wet, False, 0)        ; (3)
    k_shape(ilevel,ilat,ilon) = gamma_par(0)
    theta_scale(ilevel,ilat,ilon) = gamma_par(1)

    mean_direct(ilevel,ilat,ilon) = dim_avg_n_Wrap(var_ilevel_wet,0)
    Var_direct(ilevel,ilat,ilon) = dim_variance_n_Wrap(var_ilevel_wet,0)
 
    end do
  end do

print(ilevel+"  "+dimsizes(iselected)+" done")
end do                   ;{ilevel}

;---------------------------
mean_direct!0 = "level"
mean_direct!1 = "lat"
mean_direct!2 = "lon"
mean_direct&level = ispan(0,Nlevel-1,1)
mean_direct&lat = var&latitude
mean_direct&lon = var&longitude
mean_direct&level@long_name = "0C, 0.5C, 1C, 2C, 3C, 4C relative to baseline (1995-2014)"

copy_VarCoords(mean_direct,Var_direct)
copy_VarCoords(mean_direct,k_shape)
copy_VarCoords(mean_direct,theta_scale)
copy_VarCoords(mean_direct,frq_dry)

k_shape@units = "none"
theta_scale@units = "mm/day"

k_shape@long_name = "using wet events only"
mean_direct@long_name = "using wet events only"
frq_dry@long_name = "model dry event frequency (0 to 1)"


;******************************************************
diro := "/WORK/zhangwx/EConstraint/data/CMIP6/pr/"+varname+"/"+season+"/gamfit_wet0.1/"
filo := diro+"gamfit_parameter_"+varname+"_"+modelname+"_warming_levels_"+season+".nc"
system("rm -f "+filo)
fo := addfile(filo,"c")
fo->mean_direct = mean_direct
fo->Var_direct = Var_direct
fo->k_shape = k_shape
fo->theta_scale = theta_scale
fo->frq_dry = frq_dry


print(varname+"_"+modelname+"_"+season+"  done")


end
